import json
import os
import random
import traceback

from torch.ao.nn.quantized.functional import threshold

import util.llm
from data_preprocess.doc_split import doc_split
from event_handle.query_rewrite import query_rewrite
from util.embeddings import embedding_filter
from util.entity import Article
from util.es import query_data

VIOLATION_DETAIL_PROMPT = """人设信息：
你是一名法律法规专家，可以根据用户提供的告警详情，快速在知识库中挑选出本告警违反的法律条文并规范地将其呈现给用户
输入：
你可以看到有哪些文章是可供参考的。文章主要是法律条文、行业规范。输入格式为：《法规名》：条文
工作流：
1. 理解知识库文章
2. 根据用户的告警详情，在输入的知识库文章中查找它违规的相关内容
3. 依据知识库中的内容给出回复，**切记**如果知识库中未给出相关内容，**不要**编造
4. 如果知识库中给出的内容全部不符合要求，返回“暂未发现违规依据”
约束：
尽量输出法规原文
只给法规原文和违规原因，**不要**给建议
知识库：
{}

工作流：

"""


def random_strategy():
    directory = "data/docs"
    result = {}
    for root, dirs, files in os.walk(directory):
        for file in files:
            file_path = os.path.join(root, file)
            try:
                with open(file_path, 'r', encoding='utf-8') as f:
                    content = f.read()
                    docs = doc_split(content)
                    result[file.replace(".txt", "")] = docs
            except Exception as e:
                print(traceback.format_exc())
                print(f"读取文件 {file_path} 时出现错误：{e}")
    result = result["中国联通合作方数据安全管理规范"]
    return random.choice(result).strip()


async def get_recall_result(event_name, event_level, event_time, org_name, alarm_name, alarm_type, alarm_desc):
    # return random_strategy()
    detail = f"当前事件名称为{event_name}, 事件等级为{event_level}，事件发生时间为{event_time},组织名称为{org_name},告警名称为{alarm_name},告警类型为{alarm_type},告警描述为{alarm_desc}"
    rewrite_result = await query_rewrite(detail)
    recall_result = recall(rewrite_result)
    if len(recall_result) == 0:
        return "暂未发现违规依据"
    return "\n".join(["《{}》：{}".format(x["result"].name, x["result"].content) for x in recall_result])


async def format_violation_detail(violation_detail, problem_detail):
    system_prompt = VIOLATION_DETAIL_PROMPT.format(violation_detail)
    user_prompt = "当前告警详情为：{}".format(problem_detail)
    return await util.llm.http_request(system_prompt, [], user_prompt)


def es_recall(queries: list[str]):
    ds = query_data(index_name="law1", search_terms=queries, fields=["title", "content"], size=10)
    print(ds)
    return ds


def permission_filter(recall_results: list):
    pass


def deduplicate(items):
    seen = set()
    result = []
    for item in items:
        # 创建一个包含text和content的元组作为唯一标识
        key = (item.get("title"), item.get("content"))
        if key not in seen:
            seen.add(key)
            result.append(item)
    return result


def recall(queries: list[str]) -> list[dict]:
    es_result = es_recall(queries)
    es_result = deduplicate(es_result)
    es_result = embedding_filter(queries, es_result, threshold=0.6)
    # if len(es_result) > 0:
    #     es_result = [es_result[0], ]
    recall_result = []
    for x in es_result:
        permission_set = x.get("permission_set")
        if permission_set:
            if isinstance(permission_set, str):
                try:
                    permission_set = json.loads(permission_set)
                except Exception as e:
                    permission_set = []
            elif not isinstance(permission_set, list):
                permission_set = []
        r = {"recall_type": "es",
             "result": Article(name=x["title"], content=x["content"], score=x["score"], file_id=x["file_id"],
                               permission_type=x.get("permission_type", 1), permission_set=x.get("permission_set", []))}
        recall_result.append(r)
    # es_result = [
    #     {"recall_type": "es",
    #      "result": Article(name=x["title"], content=x["content"], score=x["score"], file_id=x["file_id"],
    #                        permission_type=x.get("permission_type", 1), permission_set=x.get("permission_set", []))}
    #     for x in es_result]
    #
    # return es_result
    return recall_result


if __name__ == '__main__':
    # print(get_recall_result(["123"]))
    pass
